from flask import request, g
from sqlalchemy import select

from config import server_config
from config.const import constant
from functools import wraps

from database.db_manager import db_session
from database.user_model import users, DBCustomQuery
from forword.custom_error.request_errors import RequestError
from forword.custom_error.token_errors import TokenExpireError

def process_token(_request):
    """process token"""
    access_code = request.headers.get("token")
    if not access_code:
        # resp = (
        #     {
        #         "error": "请在header里Authorization字段中提供有效的登陆凭证"
        #     },
        #     401,
        # )
        return None

    return server_config.token_util.parse_token(access_code)


def token_require(func):
    """
    如果某个入口函数需要token，需要在文件头加入这个装饰器
    """

    @wraps(func)
    def inner_wrapper(*args, **kwargs):
        """
        处理用户编号
        """
        user_id = process_token(request)
        user = DBCustomQuery(users).where(users.c.id == user_id).fetchone()

        if not user:
            raise RequestError("用户不存在")
        g.user_id = user_id
        return func(*args, **kwargs)

    return inner_wrapper


def admin_token_require(func):
    """
    如果某个入口函数需要token，需要在文件头加入这个装饰器
    """

    @wraps(func)
    def inner_wrapper(*args, **kwargs):
        """
        处理用户编号
        """
        user_id = process_token(request)
        user = db_session().execute(
            select(users).where(users.c.id == user_id)).fetchone()
        if "admin" not in user.roles.split(","):
            raise TokenExpireError("invalid permission")
        g.user_id = user_id
        return func(*args, **kwargs)

    return inner_wrapper

